import random
import numpy as np
from pathlib import Path
import sys
sys.path.append(str(Path(__file__).parent.parent))
from consts import ACTION
import copy

class RandomAgent:
    def get_action(self, obs):
        if any(obs['obs']['action_mask']):
            probs = copy.deepcopy(obs['obs']['action_mask'])
            # if probs[1:].any():
            #     probs[0] = 0
            if probs[:-1].any():
                probs[-1] = 0            
            probs = probs/sum(probs)
            act = np.random.choice(np.arange(6), p = probs)
        else:
            act = 0
        return [np.eye(len(ACTION))[act].astype(int).tolist()]